import os
import shutil
import argparse

# 文件分类配置（可根据需要修改）
FILE_CATEGORIES = {
    "Images": [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".svg"],
    "Documents": [".pdf", ".docx", ".doc", ".xlsx", ".pptx", ".txt", ".md"],
    "Audio": [".mp3", ".wav", ".aac", ".flac", ".m4a"],
    "Videos": [".mp4", ".mov", ".avi", ".mkv", ".flv", ".webm"],
    "Archives": [".zip", ".rar", ".7z", ".tar", ".gz"],
    "Programs": [".exe", ".msi", ".dmg", ".pkg", ".sh"],
    "Code": [".py", ".js", ".html", ".css", ".java", ".cpp"],
    "Others": []
}

def organize_files(src_dir, recursive=False):
    """
    整理指定目录中的文件
    :param src_dir: 要整理的源目录
    :param recursive: 是否递归处理子目录
    """
    processed_folders = set()

    for root, dirs, files in os.walk(src_dir):
        # 跳过已经处理过的分类文件夹
        dirs[:] = [d for d in dirs if os.path.join(root, d) not in processed_folders]

        for filename in files:
            # 跳过程序文件本身
            if filename == 'file_manage.py':
                continue
            
            # 跳过隐藏文件
            if filename.startswith('.'):
                continue

            file_path = os.path.join(root, filename)
            file_ext = os.path.splitext(filename)[1].lower()

            # 确定文件分类
            dest_category = "Others"
            for category, extensions in FILE_CATEGORIES.items():
                if file_ext in extensions:
                    dest_category = category
                    break

            # 创建目标目录
            dest_dir = os.path.join(src_dir, dest_category) if not recursive else os.path.join(root, dest_category)
            os.makedirs(dest_dir, exist_ok=True)
            processed_folders.add(dest_dir)

            # 处理同名文件
            dest_path = os.path.join(dest_dir, filename)
            if os.path.exists(dest_path):
                base, ext = os.path.splitext(filename)
                counter = 1
                while os.path.exists(dest_path):
                    new_name = f"{base}_{counter}{ext}"
                    dest_path = os.path.join(dest_dir, new_name)
                    counter += 1

            try:
                shutil.move(file_path, dest_path)
                print(f"Moved '{filename}' to '{dest_dir}'")
            except Exception as e:
                print(f"Error moving '{filename}': {str(e)}")

        # 如果不是递归模式，只处理顶层目录
        if not recursive:
            break

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="自动化文件整理工具")
    parser.add_argument("src_dir", nargs="?", default=".", help="要整理的目录（默认为当前目录）")
    parser.add_argument("-r", "--recursive", action="store_true", help="递归整理子目录")
    args = parser.parse_args()

    if not os.path.isdir(args.src_dir):
        print(f"错误：目录 '{args.src_dir}' 不存在")
        exit(1)

    print(f"开始整理目录: {os.path.abspath(args.src_dir)}")
    organize_files(args.src_dir, args.recursive)
    print("文件整理完成！")